# -*- coding: utf-8 -*-

"""
测试嵌入模型功能的简单脚本
"""

import sys
from db_utils import Database, VectorStore


def test_embedding():
    print("=== 测试嵌入模型功能 ===")

    try:
        # 初始化数据库连接
        db = Database()
        session = db.connect()
        print("✅ 数据库连接成功")

        # 初始化向量存储
        vector_store = VectorStore(db)
        print("✅ 向量存储初始化成功")

        # 测试文本
        test_text = "这是一段测试文本，用于验证嵌入模型是否正常工作。"
        print(f"\n测试文本: {test_text}")

        # 获取嵌入向量
        print("\n[正在获取嵌入向量...]")
        embedding = vector_store.get_text_embedding(test_text)

        # 验证嵌入向量
        if embedding and isinstance(embedding, list):
            print(f"✅ 嵌入向量获取成功")
            print(f"嵌入向量维度: {len(embedding)}")
            print(f"嵌入向量前5个值: {embedding[:5]}")
        else:
            print("❌ 嵌入向量获取失败")
            return False

        print("\n=== 测试成功完成 ===")
        return True

    except Exception as e:
        print(f"❌ 测试过程中出错: {str(e)}")
        import traceback

        traceback.print_exc()
        return False
    finally:
        # 确保关闭数据库连接
        if "db" in locals():
            db.disconnect(session)


if __name__ == "__main__":
    success = test_embedding()
    sys.exit(0 if success else 1)
